from scapy.all import *
from scapy.layers.inet import IP, TCP, UDP
from scapy.layers.inet6 import IPv6
from scapy.layers.l2 import Ether, ARP
from scapy.layers.ppp import PPP, HDLC


# 自定义Cisco HDLC头部
class CiscoHDLC(Packet):
    name = "Cisco HDLC"
    fields_desc = [
        XByteField("address", 0x8F),  # 默认地址字段
        ByteField("control", 0x00),  # 默认控制字段
        XShortField("protocol", 0x0800)  # 默认协议类型字段 (IPv4)
    ]


# 绑定CiscoHDLC与Raw层
# bind_layers(CiscoHDLC, Raw)
bind_layers(CiscoHDLC, IP, protocol=0x0800)

# 注册DLT类型
# 这里104对应DLT_C_HDLC，112对应DLT_HDLC
conf.l2types.register(104, CiscoHDLC)  # 注册DLT_C_HDLC
conf.l2types.register(112, CiscoHDLC)  # 注册DLT_HDLC


class Frame:
    def __init__(self, mac_layer=None, ip_layer=None, ipv6_layer=None, tcp_layer=None, udp_layer=None, arp_layer=None,
                 ppp_layer=None, hdlc_layer=None, raw=None):
        self.mac_layer = mac_layer if mac_layer is not None else {}
        self.ip_layer = ip_layer if ip_layer is not None else {}
        self.ipv6_layer = ipv6_layer if ipv6_layer is not None else {}  # IPv6层
        self.tcp_layer = tcp_layer if tcp_layer is not None else {}
        self.udp_layer = udp_layer if udp_layer is not None else {}
        self.arp_layer = arp_layer if arp_layer is not None else {}
        self.ppp_layer = ppp_layer if ppp_layer is not None else {}
        self.hdlc_layer = hdlc_layer if hdlc_layer is not None else {}
        self.raw = raw

    def add_mac_layer(self, dst, src, type):
        self.mac_layer['DestinationMACAddress'] = dst
        self.mac_layer['SourceMACAddress'] = src
        self.mac_layer['Type'] = f"0x{type:04x}"

    def add_ip_layer(self, version, ihl, tos, len, id, flags, flag, ttl, proto,
                     chksum, src, dst):
        self.ip_layer['Version'] = version
        self.ip_layer['HeaderLength'] = ihl * 4
        self.ip_layer['DifferentiatedServicesField'] = f"0x{tos:02x}"
        self.ip_layer['TotalLength'] = len
        self.ip_layer['Identification'] = id
        self.ip_layer['Flags'] = flags
        self.ip_layer['FragmentOffset'] = flag
        self.ip_layer['TimeToLive'] = ttl
        self.ip_layer['Protocol'] = proto
        self.ip_layer['HeaderChecksum'] = f"0x{chksum:04x}"
        self.ip_layer['SourceIPAddress'] = src
        self.ip_layer['DestinationIPAddress'] = dst

    def add_ipv6_layer(self, version, traffic_class, flow_label, payload_length, next_header, hop_limit, src, dst):
        self.ipv6_layer['Version'] = version
        self.ipv6_layer['TrafficClass'] = traffic_class
        self.ipv6_layer['FlowLabel'] = flow_label
        self.ipv6_layer['PayloadLength'] = payload_length
        self.ipv6_layer['NextHeader'] = next_header
        self.ipv6_layer['HopLimit'] = hop_limit
        self.ipv6_layer['SourceIPv6Address'] = src
        self.ipv6_layer['DestinationIPv6Address'] = dst

    def add_tcp_layer(self, sport, dport, seq, ack, dataofs, reserved, flags, window, chksum, urgptr):
        self.tcp_layer['SourcePort'] = sport
        self.tcp_layer['DestinationPort'] = dport
        self.tcp_layer['SequenceNumber'] = seq
        self.tcp_layer['Acknowledgment'] = ack
        self.tcp_layer['HeaderLength'] = dataofs * 4
        self.tcp_layer['Reserved'] = reserved
        self.tcp_layer['Flags'] = flags
        self.tcp_layer['Window'] = window
        self.tcp_layer['Checksum'] = f"0x{chksum:04x}"
        self.tcp_layer['UrgentPointer'] = urgptr

    def add_udp_layer(self, sport, dport, len, chksum):
        self.udp_layer['SourcePort'] = sport
        self.udp_layer['DestinationPort'] = dport
        self.udp_layer['Length'] = len
        self.udp_layer['Checksum'] = chksum

    def add_arp_layer(self, hw_type, proto_type, hw_size, proto_size, op, src_mac, src_ip, dst_mac, dst_ip):
        self.arp_layer['HardwareType'] = hw_type
        self.arp_layer['ProtocolType'] = proto_type
        self.arp_layer['HardwareSize'] = hw_size
        self.arp_layer['ProtocolSize'] = proto_size
        self.arp_layer['Operation'] = op
        self.arp_layer['SourceMACAddress'] = src_mac
        self.arp_layer['SourceIPAddress'] = src_ip
        self.arp_layer['DestinationMACAddress'] = dst_mac
        self.arp_layer['DestinationIPAddress'] = dst_ip

    def add_ppp_hdlc_like(self, address, control):
        self.ppp_layer['Address'] = f"0x{address:02x}"
        self.ppp_layer['Control'] = f"0x{control:02x}"

    def add_ppp_proto(self, proto):
        self.ppp_layer['Protocol'] = f"0x{proto:04x}"

    def add_hdlc_layer(self, address, control, protocol):
        self.hdlc_layer['Address'] = f"0x{address:02x}"
        self.hdlc_layer['Control'] = f"0x{control:02x}"
        self.hdlc_layer['Protocol'] = f"0x{protocol:04x}"

    def __repr__(self):
        # 仅显示非空字段
        repr_str = "Frame("
        if self.mac_layer:
            repr_str += f"Mac_layer={self.mac_layer}, "
        if self.hdlc_layer:
            repr_str += f"HDLC_layer={self.hdlc_layer}, "
        if self.ppp_layer:
            repr_str += f"PPP_layer={self.ppp_layer}, "
        if self.ip_layer:
            repr_str += f"IP_layer={self.ip_layer}, "
        if self.ipv6_layer:
            repr_str += f"IPv6_layer={self.ipv6_layer}, "
        if self.arp_layer:
            repr_str += f"ARP_layer={self.arp_layer}, "
        if self.tcp_layer:
            repr_str += f"TCP_layer={self.tcp_layer}, "
        if self.udp_layer:
            repr_str += f"UDP_layer={self.udp_layer}, "
        repr_str += f"raw={self.raw})"
        return repr_str


# 解析pcap文件函数
def parse_protocol(pcap_file):
    frames = []  # 这个是数据包的列表
    with PcapReader(pcap_file) as reader:
        for data in reader:
            frame = Frame()  # 创建一个新的Frame 实例
            # 解析MAC层
            if Ether in data:
                eth = data[Ether]
                frame.add_mac_layer(eth.dst, eth.src, eth.type)

            # 解析HDLC层，假设HDLC层是Cisco HDLC类型
            if CiscoHDLC in data:
                hdlc = data[CiscoHDLC]
                frame.add_hdlc_layer(
                    address=hdlc.address,
                    control=hdlc.control,
                    protocol=hdlc.protocol
                )

            # 对于scapy的源码库中会把ppp-hdlc的头部认成HDLC
            if HDLC in data:
                hdlc = data[HDLC]
                frame.add_ppp_hdlc_like(
                    address=hdlc.address,
                    control=hdlc.control
                )
                # 检查HDLC中是否包含PPP层
                if PPP in hdlc:
                    ppp = hdlc[PPP]
                    frame.add_ppp_proto(
                        proto=ppp.proto
                    )

            # 直接解析PPP层（如果未被HDLC封装）
            elif PPP in data:
                ppp = data[PPP]
                frame.add_ppp_proto(
                    proto=ppp.proto
                )

            # 解析Ip层
            if IP in data:
                ip = data[IP]
                frame.add_ip_layer(
                    version=ip.version,
                    ihl=ip.ihl,
                    tos=ip.tos,
                    len=ip.len,
                    id=ip.id,
                    flags=ip.flags,
                    flag=ip.frag,
                    ttl=ip.ttl,
                    proto=ip.proto,
                    chksum=ip.chksum,
                    src=ip.src,
                    dst=ip.dst,
                )

            # 解析IPv6层
            elif IPv6 in data:
                ipv6 = data[IPv6]
                frame.add_ipv6_layer(
                    version=ipv6.version,
                    traffic_class=ipv6.tc,
                    flow_label=ipv6.fl,
                    payload_length=ipv6.plen,
                    next_header=ipv6.nh,
                    hop_limit=ipv6.hlim,
                    src=ipv6.src,
                    dst=ipv6.dst
                )
            # 解析ARP层
            elif ARP in data:
                arp = data[ARP]
                frame.add_arp_layer(
                    hw_type=arp.hwtype,
                    proto_type=arp.ptype,
                    hw_size=arp.hwlen,
                    proto_size=arp.plen,
                    op=arp.op,
                    src_mac=arp.psrc,
                    src_ip=arp.hwsrc,
                    dst_mac=arp.pdst,
                    dst_ip=arp.hwdst
                )
            # 解析TCP层
            if TCP in data:
                tcp = data[TCP]
                frame.add_tcp_layer(
                    sport=tcp.sport,
                    dport=tcp.dport,
                    seq=tcp.seq,
                    ack=tcp.ack,
                    dataofs=tcp.dataofs,
                    reserved=tcp.reserved,
                    flags=tcp.flags,
                    window=tcp.window,
                    chksum=tcp.chksum,
                    urgptr=tcp.urgptr,
                )
                # 剩余的数据当作 raw 负载
                frame.raw = bytes(tcp.payload)

            # 解析UDP层
            elif UDP in data:
                udp = data[UDP]
                frame.add_udp_layer(
                    sport=udp.sport,
                    dport=udp.dport,
                    len=udp.len,
                    chksum=udp.chksum
                )
                # 剩余的数据当作 raw 负载
                frame.raw = bytes(udp.payload)

            # 将解析好的 frame 对象添加到 frames 列表中
            frames.append(frame)
    return frames  # 返回所有解析后的数据包列表


if __name__ == '__main__':
    frames = parse_protocol("D:\\PycharmProjects\\Protocol_parse\\pcap_file\\HTTP\\http13.cap")
    for i in frames:
        print(i)
